package com.ververica.cdc.guass.source.hybrid;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.api.connector.source.*;
import org.apache.flink.core.io.SimpleVersionedSerializer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.*;

public class HybridSourceEnumerator implements SplitEnumerator<HybridSourceSplit, HybridSourceEnumeratorState> {
    private static final Logger LOG = LoggerFactory.getLogger(HybridSourceEnumerator.class);

    private final SplitEnumeratorContext<HybridSourceSplit> context;
    private final List<HybridSourceSplit> pendingSplits;
    private int currentSourceIndex;
    //每个subTask的split分配情况
    private final Map<Integer, List<HybridSourceSplit>> assignedSplits;
    private final SimpleVersionedSerializer<HybridSourceSplit> splitSerializer;

    public HybridSourceEnumerator(
            SplitEnumeratorContext<HybridSourceSplit> context,
            List<HybridSourceSplit> initialSplits,
            int initialSourceIndex,
            HybridSourceEnumeratorState restoredState,
            SimpleVersionedSerializer<HybridSourceSplit> splitSerializer) {
        this.context = context;
        this.pendingSplits = new ArrayList<>(initialSplits);
        this.currentSourceIndex = initialSourceIndex;
        this.splitSerializer = splitSerializer;
        this.assignedSplits = new HashMap<>();
        if (restoredState != null) {
            restoreState(restoredState);
        }
    }

    @Override
    public void start() {
        LOG.info("Starting HybridSourceEnumerator with sourceIndex={}", currentSourceIndex);
    }

    @Override
    public void handleSplitRequest(int subtaskId, String requesterHostname) {
        LOG.debug("Received split request from subtask {} on host {}", subtaskId, requesterHostname);
        assignSplitsToSubtask(subtaskId);
    }

    @Override
    public void addReader(int subtaskId) {
        LOG.info("Adding reader for subtask {}", subtaskId);
        assignSplitsToSubtask(subtaskId);
    }

    private void assignSplitsToSubtask(int subtaskId) {
        // 获取当前已分配的所有 splits 和当前待分配的 splits
        List<HybridSourceSplit> allAssignedSplits = new ArrayList<>();
        for (List<HybridSourceSplit> splits : assignedSplits.values()) {
            allAssignedSplits.addAll(splits);
        }

        // 获取待分配的 splits
        List<HybridSourceSplit> allSplits = new ArrayList<>(pendingSplits);
        allSplits.removeAll(allAssignedSplits); // 从待分配的 splits 中移除已分配的 splits

        // 获取当前并行度
        int totalSubtasks = context.registeredReaders().size();
        int splitsPerSubtask = allSplits.size() / totalSubtasks;
        int remainingSplits = allSplits.size() % totalSubtasks;

        int splitIndex = 0;
        List<HybridSourceSplit> splitsForThisSubtask = new ArrayList<>();
        //btaskId < remainingSplits：这个条件判断当前的 subtaskId 是否小于 remainingSplits。
        //如果是，则说明当前子任务应该获得额外的一个 split， /因为总的 splits 数量不能均匀分配，
        //前 remainingSplits 个子任务需要多一个 split。
        for (int i = 0; i < splitsPerSubtask + (subtaskId < remainingSplits ? 1 : 0); i++) {
            if (splitIndex < allSplits.size()) {
                splitsForThisSubtask.add(allSplits.get(splitIndex));
                splitIndex++;
            }
        }

        if (!splitsForThisSubtask.isEmpty()) {
            context.assignSplits(new SplitsAssignment<>(Collections.singletonMap(subtaskId, splitsForThisSubtask)));
            assignedSplits.put(subtaskId, splitsForThisSubtask);
            LOG.info("Assigned {} splits to subtask {}", splitsForThisSubtask.size(), subtaskId);
        } else {
            context.signalNoMoreSplits(subtaskId);
            LOG.info("Signaled no more splits for subtask {}", subtaskId);
        }
    }

    @Override
    public void addSplitsBack(List<HybridSourceSplit> splits, int subtaskId) {
        LOG.info("Adding {} splits back from subtask {}", splits.size(), subtaskId);
        pendingSplits.addAll(splits);
        assignedSplits.remove(subtaskId);
        assignSplitsToSubtask(subtaskId);
    }

    @Override
    public HybridSourceEnumeratorState snapshotState(long checkpointId) throws Exception {
        LOG.debug("Snapshotting state for checkpoint {}", checkpointId);
        byte[] serializedSplits = serializePendingSplits();
        return new HybridSourceEnumeratorState(currentSourceIndex, serializedSplits, splitSerializer.getVersion());
    }

    @Override
    public void close() throws IOException {
        LOG.info("Closing HybridSourceEnumerator");
    }

    private void restoreState(HybridSourceEnumeratorState state) {
        this.currentSourceIndex = state.getCurrentSourceIndex();
        try {
            List<HybridSourceSplit> restoredSplits = deserializePendingSplits(state.getWrappedState(), state.getWrappedStateSerializerVersion());
            this.pendingSplits.clear();
            this.pendingSplits.addAll(restoredSplits);
        } catch (IOException e) {
            throw new RuntimeException("Failed to restore enumerator state", e);
        }
        LOG.info("Restored enumerator state: sourceIndex={}, pendingSplits={}", currentSourceIndex, pendingSplits.size());
    }

    private byte[] serializePendingSplits() throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        return mapper.writeValueAsBytes(pendingSplits);
    }

    private List<HybridSourceSplit> deserializePendingSplits(byte[] serializedSplits, int serializerVersion) throws IOException {
        ObjectMapper mapper = new ObjectMapper();
        return mapper.readValue(serializedSplits, new TypeReference<List<HybridSourceSplit>>() {});
    }

    @Override
    public void handleSourceEvent(int subtaskId, SourceEvent sourceEvent) {
        // 处理来自 SourceReader 的事件，如果需要的话
        LOG.debug("Received source event from subtask {}: {}", subtaskId, sourceEvent);
    }
}
